import chromadb
from chromadb.config import Settings
import yaml
import json
import os
import time
import logging
from typing import List, Dict, Any, Optional
from sentence_transformers import SentenceTransformer
import asyncio

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 设置HuggingFace镜像地址，解决连接问题
hf_mirror = 'https://hf-mirror.com'
os.environ['HF_ENDPOINT'] = hf_mirror
logger.info(f"已设置HuggingFace镜像地址: {hf_mirror}")

# 设置下载超时时间
hf_timeout = '120'
os.environ['HF_HUB_DOWNLOAD_TIMEOUT'] = hf_timeout
logger.info(f"已设置HuggingFace下载超时: {hf_timeout}秒")

from app.core.config import settings

# 全局变量
client = None
collection = None
embedding_model = None

# 初始化Chroma数据库
async def init_db():
    global client, collection, embedding_model
    
    logger.info("开始数据库异步初始化...")
    
    # 尝试同步初始化
    await sync_init_db()
    
    # 验证初始化结果
    if collection is not None and embedding_model is not None:
        logger.info("数据库初始化完成并验证成功")
    else:
        logger.warning("数据库初始化完成但验证失败:")
        logger.warning(f"  - client: {client is not None}")
        logger.warning(f"  - collection: {collection is not None}")
        logger.warning(f"  - embedding_model: {embedding_model is not None}")

# 同步初始化数据库
async def sync_init_db():
    global client, collection, embedding_model
    
    logger.info("开始同步初始化数据库...")
    
    # 使用异步方式加载嵌入模型（在默认执行器中运行同步函数）
    logger.info(f"准备加载嵌入模型: {settings.TEXT_EMBEDDING_MODEL}")
    model = await asyncio.to_thread(load_embedding_model_with_retry, settings.TEXT_EMBEDDING_MODEL)
    
    # 如果主模型加载失败，尝试备选模型
    if model is None:
        logger.warning("主嵌入模型未成功加载，尝试使用备选模型")
        model = await asyncio.to_thread(load_embedding_model_with_retry, "all-MiniLM-L6-v2")
    
    if model is None:
        logger.error("无法加载任何嵌入模型")
        return
    
    # 明确设置全局变量
    logger.info("嵌入模型加载成功，更新全局变量")
    embedding_model = model
    logger.info(f"嵌入模型全局变量设置完成: {embedding_model is not None}")
    
    # 初始化Chroma客户端
    try:
        # 通过环境变量禁用遥测功能
        if hasattr(settings, 'CHROMA_DISABLE_TELEMETRY') and settings.CHROMA_DISABLE_TELEMETRY:
            os.environ['CHROMA_TELEMETRY_ENABLED'] = 'false'
            logger.info("已通过环境变量禁用ChromaDB遥测功能")
        
        # 使用异步方式初始化客户端
        logger.info(f"准备初始化Chroma客户端，持久化路径: {settings.CHROMA_DB_PATH}")
        _client = await asyncio.to_thread(
            chromadb.Client,
            Settings(
                persist_directory=settings.CHROMA_DB_PATH,
                is_persistent=True
            )
        )
        
        # 明确设置全局变量
        client = _client
        logger.info(f"Chroma客户端全局变量设置完成: {client is not None}")
        
        # 创建或获取集合
        logger.info("准备创建或获取集合: sql_knowledge")
        _collection = await asyncio.to_thread(
            client.get_or_create_collection,
            name="sql_knowledge",
            metadata={"description": "Text2SQL知识集合，包含数据库模式和查询对"}
        )
        
        # 明确设置全局变量
        collection = _collection
        logger.info(f"集合全局变量设置完成: {collection is not None}")
        
        # 检查集合是否为空，如果为空则加载数据
        count = await asyncio.to_thread(collection.count)
        if count == 0:
            try:
                await load_initial_data()
            except Exception as e:
                logger.error(f"加载初始数据失败: {str(e)}")
                # 即使加载数据失败，服务仍然可以启动，只是功能会受限
        else:
            logger.info(f"集合已包含 {count} 条文档，跳过数据加载")
            
    except Exception as e:
        logger.error(f"数据库初始化失败: {str(e)}")
        # 确保全局变量不会保留部分初始化的状态
        client = None
        collection = None

# 加载初始数据（数据库模式和查询对）
async def load_initial_data():
    try:
        # 使用相对路径代替绝对路径
        base_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), 'sakila')
        
        # 加载数据库描述
        db_desc_path = os.path.join(base_path, 'db_description.yaml')
        if not os.path.exists(db_desc_path):
            logger.error(f"数据库描述文件不存在: {db_desc_path}")
            return
        
        with open(db_desc_path, 'r', encoding='utf-8') as f:
            db_description = yaml.safe_load(f)
        
        # 加载DDL语句
        ddl_path = os.path.join(base_path, 'ddl_statements.yaml')
        if not os.path.exists(ddl_path):
            logger.error(f"DDL语句文件不存在: {ddl_path}")
            return
        
        with open(ddl_path, 'r', encoding='utf-8') as f:
            ddl_statements = yaml.safe_load(f)
        
        # 加载问题-答案对
        q2sql_path = os.path.join(base_path, 'q2sql_pairs.json')
        if not os.path.exists(q2sql_path):
            logger.error(f"问题-答案对文件不存在: {q2sql_path}")
            return
        
        with open(q2sql_path, 'r', encoding='utf-8') as f:
            q2sql_pairs = json.load(f)
        
        # 准备要插入的数据
        documents = []
        metadatas = []
        ids = []
        
        # 添加数据库描述文档
        for table_name, columns in db_description.items():
            doc = f"表名：{table_name}\n列信息：\n"
            for col_name, col_desc in columns.items():
                doc += f"- {col_name}: {col_desc}\n"
            
            documents.append(doc)
            metadatas.append({"type": "table_description", "table_name": table_name})
            ids.append(f"table_{table_name}")
        
        # 添加DDL语句文档
        for table_name, ddl in ddl_statements.items():
            documents.append(ddl)
            metadatas.append({"type": "ddl_statement", "table_name": table_name})
            ids.append(f"ddl_{table_name}")
        
        # 添加问题-答案对文档
        for i, pair in enumerate(q2sql_pairs):
            documents.append(f"问题：{pair['question']}\nSQL：{pair['sql']}")
            metadatas.append({"type": "q2sql_pair", "question": pair['question'], "sql": pair['sql']})
            ids.append(f"q2sql_{i}")
        
        # 生成嵌入向量
        if embedding_model is None:
            logger.error("嵌入模型未初始化，无法生成嵌入向量")
            return
            
        embeddings = embedding_model.encode(documents).tolist()
        
        # 插入数据到集合
        collection.add(
            documents=documents,
            embeddings=embeddings,
            metadatas=metadatas,
            ids=ids
        )
        
        logger.info(f"成功加载 {len(documents)} 条初始数据到集合")
    except Exception as e:
        logger.error(f"加载初始数据时出错: {str(e)}")
        raise

# 加载嵌入模型，带重试机制
def load_embedding_model_with_retry(model_name: str, max_retries: int = 3, retry_delay: int = 5) -> Optional[SentenceTransformer]:
    """加载嵌入模型，带重试机制和错误处理"""
    retries = 0
    
    while retries < max_retries:
        try:
            logger.info(f"正在加载嵌入模型: {model_name} (尝试 {retries + 1}/{max_retries})")
            
            # 设置本地缓存目录
            cache_dir = os.path.join(settings.CHROMA_DB_PATH, "model_cache")
            os.makedirs(cache_dir, exist_ok=True)
            
            model = SentenceTransformer(
                model_name,
                cache_folder=cache_dir,
                use_auth_token=False  # 不使用认证令牌
            )
            
            logger.info(f"成功加载嵌入模型: {model_name}")
            return model
        except Exception as e:
            retries += 1
            logger.warning(f"加载模型失败 (尝试 {retries}/{max_retries}): {str(e)}")
            
            if retries < max_retries:
                logger.info(f"{retry_delay}秒后重试...")
                time.sleep(retry_delay)
            else:
                logger.error(f"所有重试均失败，无法加载模型: {model_name}")
                # 返回None，让应用程序可以在没有模型的情况下继续运行（虽然功能会受限）
                return None

# 查询相关文档
async def query_documents(query_text: str, n_results: int = 5) -> List[Dict[str, Any]]:
    global collection, embedding_model
    
    try:
        # 检查并确保数据库已初始化
        if collection is None or embedding_model is None:
            logger.warning("集合或嵌入模型未初始化，尝试同步初始化")
            await sync_init_db()
            
            # 再次检查
            if collection is None or embedding_model is None:
                logger.error("数据库初始化失败，无法执行查询")
                return []
        
        # 生成查询嵌入向量
        query_embedding = await asyncio.to_thread(embedding_model.encode, [query_text])
        query_embedding = query_embedding.tolist()
        
        # 查询相关文档
        results = await asyncio.to_thread(
            collection.query,
            query_embeddings=query_embedding,
            n_results=n_results,
            include=['documents', 'metadatas', 'distances']
        )
        
        # 格式化结果
        formatted_results = []
        if results['ids'] and results['ids'][0]:
            for i in range(len(results['ids'][0])):
                formatted_results.append({
                    'id': results['ids'][0][i],
                    'document': results['documents'][0][i],
                    'metadata': results['metadatas'][0][i],
                    'distance': results['distances'][0][i]
                })
        
        return formatted_results
    except Exception as e:
        logger.error(f"查询文档时出错: {str(e)}")
        return []